from flask import Flask, request, session, jsonify
import difflib, json, PyPDF2, importlib
from pygments import highlight
from pygments.lexers import PythonLexer, get_lexer_by_name
from pygments.formatters import HtmlFormatter

app = Flask(__name__)
app.secret_key = 'innovation_diff_tool_secret'

try:
    docx_module = importlib.import_module('docx')
    Document = getattr(docx_module, 'Document', None)
    if Document is None:
        raise ImportError
    docx_supported = True
except ImportError:
    Document = None
    docx_supported = False

formatter = HtmlFormatter(cssclass='highlight', nowrap=True)

CODE_LANG_MAP = {
    'py': 'python', 'js': 'javascript', 'cpp': 'cpp',
    'java': 'java', 'ipynb': 'python'
}


def escape_html(t):
    return t.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')


def extract_code_cells(nb):
    return [line for c in nb.get('cells', []) if c.get('cell_type') == 'code' for line in c.get('source', []) if
            line.strip()]


def smart_extract(fs):
    ext = fs.filename.rsplit('.', 1)[-1].lower()
    if ext == 'ipynb':
        try:
            code = extract_code_cells(json.load(fs.stream))
            return code, 'Notebook', 'python'
        except:
            pass
    if ext in CODE_LANG_MAP:
        fs.stream.seek(0)
        return fs.stream.read().decode('utf-8', 'ignore').splitlines(), 'Text', CODE_LANG_MAP[ext]
    if ext == 'docx':
        if docx_supported and Document:
            return [p.text for p in Document(fs.stream).paragraphs], 'Text', 'word'
        fs.stream.seek(0)
        return fs.stream.read().decode('utf-8', 'ignore').splitlines(), 'Text', 'word'
    if ext == 'pdf':
        return [l for p in PyPDF2.PdfReader(fs.stream).pages for l in
                (p.extract_text() or '').splitlines()], 'Text', 'pdf'
    fs.stream.seek(0)
    return fs.stream.read().decode('utf-8', 'ignore').splitlines(), 'Text', 'text'


def generate_diff_output(old, new, diff_only, lang):
    l, r = [], []
    seq = difflib.SequenceMatcher(None, old, new)
    for tag, i1, i2, j1, j2 in seq.get_opcodes():
        if diff_only and tag == 'equal':
            continue
        for k in range(max(i2 - i1, j2 - j1)):
            left_line = old[i1 + k] if k < i2 - i1 else ''
            right_line = new[j1 + k] if k < j2 - j1 else ''
            lineno = i1 + k + 1 if k < i2 - i1 else ''
            symbol = '-' if tag == 'delete' else '+' if tag == 'insert' else '~'
            num = f"<span class='line-number'>{lineno}</span>"
            sym = f"<span class='line-symbol'>{symbol}</span>"

            left_content = highlight(left_line, get_lexer_by_name(lang), formatter) if left_line else ''
            right_content = highlight(right_line, get_lexer_by_name(lang), formatter) if right_line else ''

            left_class = 'removed-line' if tag in ('delete', 'replace') else ''
            right_class = 'added-line' if tag in ('insert', 'replace') else ''

            l.append(f"<div class='{left_class}'>{num}{sym}{left_content}</div>")
            r.append(f"<div class='{right_class}'>{num}{sym}{right_content}</div>")
    return ''.join(l), ''.join(r)

